#include <opencv2/opencv.hpp>
#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <geometry_msgs/PoseStamped.h>
#include <std_msgs/String.h>
#include <std_msgs/String.h>
#include <cv_bridge/cv_bridge.h>

#include <iostream>
#include <vector>
#include <string>
#include <stdio.h>
#include "chrono"

using namespace std;
using namespace cv;

#define IS_VIDEO 

//--------------function declaration----------------
void process(Mat& frame);
int  detectColor(Mat& frame);
void publish(ros::Publisher& pub);
void print();

//--------------global variables stored to publish----------------
struct Target
{
    int    color    = 0; // 1: red, 2: green, 3: blue 4: yellow
    int    shape    = 0; // 1: square, 2: circle, 3: triangle
    double offset_x = 0.0; // up:0~100 down:-100~0
    double offset_y = 0.0; // left:0~100 right:-100~0}
};
vector<Target> targets;

//----------------main----------------
int main(int argc, char **argv)
{
#ifdef IS_VIDEO
    ros::init(argc, argv, "vision_node");
    ros::NodeHandle nh;
    ros::Publisher vision_pub = nh.advertise<geometry_msgs::PoseStamped>("/ele_vision", 10);
    ros::Rate rate(30);

    VideoCapture cap(2);
    cap.set(CAP_PROP_FOURCC, VideoWriter::fourcc('M', 'J', 'P', 'G'));
    cap.set(CAP_PROP_FRAME_WIDTH, 640);
    cap.set(CAP_PROP_FRAME_HEIGHT, 480);
    if (!cap.isOpened())cout << "Error opening video stream or file" << endl;
    
    while (ros::ok())
    {
        Mat frame;
        cap >> frame;
#else
        Mat frame;
        frame = imread("/home/firefly/Pictures/map01.png");
        cv::resize(frame, frame, cv::Size(640, 480));
#endif
        if(frame.empty())cout << "frame is empty" << endl;
        process(frame);
        imshow("Frame", frame);

#ifdef IS_VIDEO
        publish(vision_pub);
        char c = (char)waitKey(1);
        if (c == 27)break;
    }
#else
        print();
        waitKey(0);
#endif
        
    return 0;
}

//--- -----process img then store in global varities------------
void process(Mat& frame)
{
    Mat                   gray, thresh;
    vector<vector<Point>> contours;
    vector<Vec4i>         hierarchy;
    
    cvtColor(frame, gray, COLOR_BGR2GRAY);
    Canny(gray, thresh, 69, 200);
    // threshold(gray, thresh, 100, 255, THRESH_BINARY_INV);
    findContours(thresh, contours, hierarchy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);

    imshow("thresh", thresh);
    for (size_t i = 0; i < contours.size(); i++)
    {
        double            area, bboxArea, ratio;
        int               cX, cY;
        RotatedRect       minRect;
        Target            target;
        Point2f           rect_points[4];
        Moments           M;
        Mat               framez;
        area     = contourArea(contours[i]);
        if (area < 3000)continue;
        minRect  = minAreaRect(contours[i]);
        bboxArea = minRect.size.area();
        ratio    = area / bboxArea;

        cout << "contour" << i << "area: " << area     << endl;
        cout << "rect"    << i << "area: " << bboxArea << endl;
        cout << "ratio: " << ratio                     << endl;

        //calculate center
        M = moments(contours[i]);
        
        if (M.m00 != 0) 
        {
            cX = int(M.m10 / M.m00);
            cY = int(M.m01 / M.m00);
        }else 
        {
            cX = frame.cols/2;
            cY = frame.rows/2;
        }

        //store target
        if (ratio  > 0.82) target.shape = 1; //square
        else if (ratio < 0.82 && ratio >0.70) target.shape = 2; //circle
        else if (ratio < 0.55 && ratio >0.45) target.shape = 3; //triangle
        else target.shape = 0;

        if (target.shape)
        {
            // frame.copyTo(framez);
            framez = Mat::zeros(frame.size(), frame.type());
            Mat maskROI = Mat::zeros(frame.size(), CV_8UC1);
            vector<vector<Point>> currentContour = { contours[i] };
            fillPoly(maskROI, currentContour, Scalar(255));
            bitwise_and(frame, frame, framez, maskROI);
            target.color = detectColor(framez);
        }

        target.offset_x = (frame.rows/2.0 - cY)/frame.rows;
        target.offset_y = (frame.cols/2.0 - cX)/frame.cols;

        //draw rect
        minRect.points(rect_points);
        for (int j = 0; j < 4; j++) line(frame, rect_points[j], rect_points[(j+1)%4], Scalar(0, 255, 0), 2);
        circle(frame, Point(cX, cY), 5, Scalar(255, 0, 0), -1); // cX: width cY: height

        targets.push_back(target);
    }

}

//--------------read global varities then publish----------------
void publish(ros::Publisher& pub) 
{
    for(auto target : targets)
    {   
        if(!target.shape || !target.color)continue;
        geometry_msgs::PoseStamped pose_msg;
        pose_msg.header.stamp       =   ros::Time::now();
        pose_msg.header.frame_id    =   "camera_coordinate"; 
        pose_msg.pose.position.x    =   target.color; 
        pose_msg.pose.position.y    =   target.shape;  
        pose_msg.pose.position.z    =   target.offset_x;    
        pose_msg.pose.orientation.w =   target.offset_y; 
        pub.publish(pose_msg);
        ROS_INFO("Published Red Square: [%f, %f, %f]", pose_msg.pose.position.x, pose_msg.pose.position.y, pose_msg.pose.position.z);
    }
    targets.clear();
}

//--------------print global varities----------------
void print()
{
    for(auto target : targets)
    {
        cout << "color: "    << target.color    << endl;
        cout << "shape: "    << target.shape    << endl;
        cout << "offset_x: " << target.offset_x << endl;
        cout << "offset_y: " << target.offset_y << endl;
    }
    targets.clear();
}

//--------------detect color and return num----------------
/* 1: red, 2: green, 3: blue 4: yellow */
int detectColor(Mat& frame)
{
    Mat hsv, mask_red1, mask_red2, mask_red, mask_yellow, mask_green, mask_blue;
    cvtColor(frame, hsv, COLOR_BGR2HSV);
    Scalar lower_red1  (0,   21,  22 );
    Scalar upper_red1  (11,  255, 255);
    Scalar lower_red2  (160, 43,  46 );
    Scalar upper_red2  (180, 255, 255);

    Scalar lower_yellow(26,  43,  46 );
    Scalar upper_yellow(34,  255, 255);
    Scalar lower_green (35,  43,  46 );
    Scalar upper_green (77,  255, 255);
    Scalar lower_blue  (100, 43,  46 );
    Scalar upper_blue  (124, 255, 255);

    inRange(hsv, lower_red1, upper_red1, mask_red1);
    inRange(hsv, lower_red2, upper_red2, mask_red2);
    // 合并两个红色掩模
    bitwise_or(mask_red1, mask_red2, mask_red);
    imshow("green ", mask_red);

    inRange(hsv, lower_yellow, upper_yellow, mask_yellow);
    inRange(hsv, lower_green, upper_green, mask_green);
    inRange(hsv, lower_blue, upper_blue, mask_blue);
    
    Mat gray_frame;
    cvtColor(frame, gray_frame, COLOR_BGR2GRAY);
    int totalNonZero = countNonZero(gray_frame);
    int redNonZero = countNonZero(mask_red);
    int yellowNonZero = countNonZero(mask_yellow);
    int greenNonZero = countNonZero(mask_green);
    int blueNonZero = countNonZero(mask_blue);
    cout << totalNonZero << " " << redNonZero << " " << yellowNonZero << " " << greenNonZero << " " << blueNonZero << endl;
    
    double redRatio = static_cast<double>(redNonZero) / totalNonZero;
    double yellowRatio = static_cast<double>(yellowNonZero) / totalNonZero;
    double greenRatio = static_cast<double>(greenNonZero) / totalNonZero;
    double blueRatio = static_cast<double>(blueNonZero) / totalNonZero;

    // 判断是否有颜色占比超过80%
    if (redRatio > 0.8) return 1;     // red
    if (yellowRatio > 0.8) return 4;  // yellow
    if (greenRatio > 0.8) return 2;   // green
    if (blueRatio > 0.8) return 3;    // blue

    // 如果没有任何颜色超过80%，返回0表示未检测到
    return 0;

}